{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Load and read the dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Read the dataset from hdf5 file and show the data struction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import h5py\n",
    "import json\n",
    "import os\n",
    "import random\n",
    "\n",
    "def get_all_hdf5_files(directory):\n",
    "    \"\"\"\n",
    "    Get all HDF5 files in a directory and its subdirectories.\n",
    "    \"\"\"\n",
    "    hdf5_files = []\n",
    "    for root, dirs, files in os.walk(directory):\n",
    "        for file in files:\n",
    "            if file.endswith('.hdf5'):\n",
    "                hdf5_files.append(os.path.join(root, file))\n",
    "    return hdf5_files\n",
    "\n",
    "def print_structure(group, indent=0):\n",
    "    \"\"\"\n",
    "    Print the hierarchical structure of an HDF5 group.\n",
    "    \"\"\"\n",
    "    for key in group.keys():\n",
    "        item = group[key]\n",
    "        print(\" \" * indent + key)\n",
    "        if isinstance(item, h5py.Group):\n",
    "            print_structure(item, indent + 4)\n",
    "        elif isinstance(item, h5py.Dataset):\n",
    "            print(\" \" * indent, item.shape, item.dtype)\n",
    "            if key == \"entities\" or key == \"target_entity\" or key == \"instruction\":\n",
    "                decoded = [x.decode('utf-8') for x in item]\n",
    "                print(\" \" * indent, decoded)\n",
    "        \n",
    "        elif isinstance(item, h5py.Datatype):\n",
    "            print(\" \" * (indent + 4) + \"Datatype\", json.loads())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['/media/shiduo/LENOVO_USB_HDD/dataset/VLABench/select_billiards/data_0.hdf5']\n"
     ]
    }
   ],
   "source": [
    "dataset_root = \"/media/shiduo/LENOVO_USB_HDD/dataset/VLABench/select_billiards\" \n",
    "hdf5_files = get_all_hdf5_files(dataset_root)\n",
    "print(hdf5_files[:3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "data\n",
      "    2025-02-23 20:46:40\n",
      "        instruction\n",
      "         (1,) |S38\n",
      "         ['Please put the striped_10 in any hole.']\n",
      "        meta_info\n",
      "            entities\n",
      "             (6,) |S15\n",
      "             ['billiards_table', 'striped_10', 'striped_14', 'striped_11', 'striped_12', 'solid_1']\n",
      "            target_entity\n",
      "             (1,) |S10\n",
      "             ['striped_10']\n",
      "        observation\n",
      "            depth\n",
      "             (212, 4, 480, 480) float32\n",
      "            ee_state\n",
      "             (212, 8) float32\n",
      "            point_cloud_colors\n",
      "             (212, 11905, 3) float32\n",
      "            point_cloud_points\n",
      "             (212, 11905, 3) float32\n",
      "            q_acceleration\n",
      "             (212, 7, 1) float32\n",
      "            q_state\n",
      "             (212, 7, 1) float32\n",
      "            q_velocity\n",
      "             (212, 7, 1) float32\n",
      "            rgb\n",
      "             (212, 4, 480, 480, 3) uint8\n",
      "            robot_mask\n",
      "             (212, 4, 480, 480) float32\n",
      "        trajectory\n",
      "         (212, 8) float32\n",
      "    2025-02-23 20:49:53\n",
      "        instruction\n",
      "         (1,) |S38\n",
      "         ['Please put the striped_14 in any hole.']\n",
      "        meta_info\n",
      "            entities\n",
      "             (5,) |S15\n",
      "             ['billiards_table', 'striped_14', 'solid_7', 'striped_15', 'solid_1']\n",
      "            episode_config\n",
      "             () |S1879\n",
      "            target_entity\n",
      "             (1,) |S10\n",
      "             ['striped_14']\n",
      "        observation\n",
      "            depth\n",
      "             (164, 4, 480, 480) float32\n",
      "            ee_state\n",
      "             (164, 8) float32\n",
      "            point_cloud_colors\n",
      "             (164, 11929, 3) float32\n",
      "            point_cloud_points\n",
      "             (164, 11929, 3) float32\n",
      "            q_acceleration\n",
      "             (164, 7, 1) float32\n",
      "            q_state\n",
      "             (164, 7, 1) float32\n",
      "            q_velocity\n",
      "             (164, 7, 1) float32\n",
      "            rgb\n",
      "             (164, 4, 480, 480, 3) uint8\n",
      "            robot_mask\n",
      "             (164, 4, 480, 480) float32\n",
      "        trajectory\n",
      "         (164, 8) float32\n"
     ]
    }
   ],
   "source": [
    "example_file = random.choice(hdf5_files)\n",
    "with h5py.File(example_file, 'r') as f:\n",
    "    print_structure(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Convert to tf dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the baseline methods, OpenVLA and Octo are trained using datasets in the TFDS format. In this repository, we have modified the tfds dataset builder to convert the HDF5 source format of VLABench into TFDS format. We offer [singlethread tfds builder](../VLABench/utils/rlds_builder.py) and [multithread tfds builder](../VLABench/utils/multithread_rlds_builder.py)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To convert the dataset, run\n",
    "```sh\n",
    "python scripts/convert_to_rlds.py --save_dir /your/vlabench/dataset/root --task {target_task}\n",
    "```\n",
    "to create the tfds builder in the dataset directory.\n",
    "Then, run\n",
    "```sh\n",
    "cd /your/vlabench/dataset/root/target_task\n",
    "tfds build --overwrite\n",
    "```\n",
    "to generate rlds format dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We recommend the readers to refer to https://github.com/kpertsch/rlds_dataset_builder for further guidance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Convert to Lerobot dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the baseline methods, $\\pi_0$ uses lerobot format dataset for finetining. We offer a script to convert hdf5 format dataset into lerobot format, similar to Libero dataset used in $\\pi_0$.\n",
    "\n",
    "Run \n",
    "```sh\n",
    "python scripts/convert_to_lerobot.py --dataset-name xxx --dataset-path /your/path/to/hdf5 --max-files 500 --task-list task1 task2 ... task n\n",
    "```\n",
    "to create a multi-task lerobot dataset."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vlabench",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
